#!/usr/bin/env python3
import os

from data_utils import *
from publish_utils import *
from kitti_util import *
from head_function import *
from tracklet_parser import parse_xml
from collections import deque

DATA_PATH = '/home/jony/data/kitti/RawData/2011_09_26/2011_09_26_drive_0061_sync/'
EGOCAR = np.array([[2.15, 0.9, -1.73], [2.15, -0.9, -1.73], [-1.95, -0.9, -1.73], [-1.95, 0.9, -1.73],
                   [2.15, 0.9, -0.23], [2.15, -0.9, -0.23], [-1.95, -0.9, -0.23], [-1.95, 0.9, -0.23]])

def compute_distance_from_velocity(vf, vl, vu, time_interval):
    """
    Compute the displacement vector from velocities.
    """
    vf, vl, vu = vf.iloc[0], vl.iloc[0], vu.iloc[0]
    displacement = np.array([vf * time_interval, vl * time_interval, vu * time_interval])
    return displacement

class Object():
    def __init__(self, center):
        # 初始化时确保center是3维坐标
        self.locations = deque(maxlen=20)
        self.locations.appendleft(center)

    def update(self, center, displacement, yaw_change):
        for i in range(len(self.locations)):
            x0, y0, z0 = self.locations[i]
            x1 = x0 * np.cos(yaw_change) - y0 * np.sin(yaw_change) + displacement[0]
            y1 = x0 * np.sin(yaw_change) + y0 * np.cos(yaw_change) + displacement[1]
            z1 = z0 + displacement[2]
            self.locations[i] = np.array([x1, y1, z1])

        # 更新时确保center是3维坐标
        if center is not None and len(center) == 3:
            self.locations.appendleft(center)

    def reset(self):
        self.locations = deque(maxlen=20)  # 清空轨迹位置并保留最大长度限制


if __name__ == '__main__':
    frame = 0
    rospy.init_node('kitti_node', anonymous=True)
    cam_pub = rospy.Publisher('kitti_cam', Image, queue_size=10)
    pcl_pub = rospy.Publisher('kitti_pcl', PointCloud2, queue_size=10)
    ego_pub = rospy.Publisher('kitti_ego_car', MarkerArray, queue_size=10)
    imu_pub = rospy.Publisher('kitti_imu', Imu, queue_size=10)
    gps_pub = rospy.Publisher('kitti_gps', NavSatFix, queue_size=10)
    box3d_pub = rospy.Publisher('kitti_3d', MarkerArray, queue_size=10)
    loc_pub = rospy.Publisher('kitti_loc', MarkerArray, queue_size=10)
    dist_pub = rospy.Publisher('kitti_dist', MarkerArray, queue_size=10)
    bridge = CvBridge()

    rate = rospy.Rate(10)

    tracklets = parse_xml('/home/jony/data/kitti/tracklets_0061/2011_09_26/2011_09_26_drive_0061_sync/tracklet_labels.xml')
    calib = Calibration('/home/jony/data/kitti/RawData/2011_09_26/', from_video=True)
    # df_types = get_frame_object_types(tracklets)
    tracker = {}  # track_id : Object
    prev_imu_data = None

    while not rospy.is_shutdown():
        # types = np.array(df_types[df_types.frame==frame]['type'])
        time_interval = 0.1
        centers = {}  # track_id: center
        minPQDs = [] 
        types = []
        boxes_2d = []
        corners_3d_velos = []
        for track_id, tracklet in enumerate(tracklets):
            # 如果当前帧数超过该tracklet的帧数，跳过该tracklet
            if frame >= tracklet.num_frames:
                continue
            h, w, l = tracklet.size[0], tracklet.size[1], tracklet.size[2]
            tx, ty, tz = tracklet.trans[frame]
            rz = wrap_to_pi(tracklet.rot[frame, 2])
            yaw = rz_to_ry(rz)
            type = tracklet.object_type
            types.append(type)
            # 将世界坐标系下的 tx, ty, tz 转换为相机坐标系下的 pos_x, pos_y, pos_z
            pos_x, pos_y, pos_z = convert_world_to_camera_coords(tx, ty, tz, calib)
            # 使用转换后的坐标计算3D框
            corners_2d, corners_3d_cam2 = compute_3d_box_cam2(h, w, l, pos_x, pos_y, pos_z, yaw, calib)
            corners_3d_velo = calib.project_rect_to_velo(corners_3d_cam2.T)
            minPQDs += [min_distance_cuboids(EGOCAR, corners_3d_velo)]
            centers[track_id] = np.mean(corners_3d_velo, axis=0)[:3]  # 保证中心点是3维坐标
            corners_3d_velos.append(corners_3d_velo)
            box_2d = get_2d_boxes(corners_3d_cam2)
            boxes_2d.append(box_2d)
        centers[-1] = np.array([0, 0, 0])  # 添加自身车辆的行车轨迹

        image = read_camera(os.path.join(DATA_PATH, 'image_02/data/%010d.png' % frame))
        point_cloud = read_point_cloud(os.path.join(DATA_PATH, 'velodyne_points/data/%010d.bin' % frame))
        imu_data = read_imu(os.path.join(DATA_PATH, 'oxts/data/%010d.txt' % frame))

        if prev_imu_data is None:
            for track_id in centers:
                tracker[track_id] = Object(centers[track_id])
        else:
            displacement = compute_distance_from_velocity(imu_data.vf, imu_data.vl, imu_data.vu, time_interval)
            yaw_change = imu_data['yaw'].iloc[0] - prev_imu_data['yaw'].iloc[0]
            for track_id in centers:
                if track_id in tracker:
                    tracker[track_id].update(centers[track_id], displacement, yaw_change)
                else:
                    tracker[track_id] = Object(centers[track_id])
            for track_id in tracker:
                if track_id not in centers:
                    tracker[track_id].update(None, displacement, yaw_change)

        prev_imu_data = imu_data
        publish_camera(cam_pub, bridge, image, boxes_2d, types)
        publish_point_cloud(pcl_pub, point_cloud)
        publish_ego_car(ego_pub)
        publish_imu(imu_pub, imu_data)
        publish_gps(gps_pub, imu_data)
        publish_3dbox(box3d_pub, corners_3d_velos, types)
        publish_loc(loc_pub, tracker, centers)
        publish_dist(dist_pub, minPQDs)
        rospy.loginfo('published frame %d' % frame)
        rate.sleep()
        frame += 1
        frame %= 703
        # # 添加一个判断，防止frame变量的死循环
        # if frame >= max([t.num_frames for t in tracklets]):
        #     frame = 0
        #     for track_id in tracker:
        #         tracker[track_id].reset()
